#include "Matrix.h"
#include "Encoder.h"
#include "FileReader.h"
#include "FileWriter.h"
#include "EncoderSettings.h"
#include "BitStream.h"
#include "Decoder.h"
#include <string>
#include <stdint.h>
#include <iostream>
#include "parallelSettings.h"
#include "nvToolsExt.h"
#include <omp.h>
#include "vld.h"
#include "NsightDebug.h"
#include <sstream>

int main(int argc, char *argv[])
{
	nvtxNameCategoryA(1, "Encode PFrame");
	nvtxRangePushA("Main");
    FileReader filereader;
    FileWriter filewriter;
    EncoderSettings config;
    QuantizationMatrix* qm;
    string filename;
    Frame *buff;
    bool success=false;

	omp_set_num_threads(NUM_THREADS);

    if(argc != 2)
    {
        printf("Invalid commandline format, should be \"encoder.exe <configfile>\", press any key to continue\n");
        getchar();
        return 1;
    }
    else
    {
        printf("Encoder starting...\n");

        printf("Reading configuration in %s...\n", argv[1]);

		filename.clear();
        filename.append(argv[1]);

        config = filereader.readEncoderSettings(filename);
        qm = filereader.readMatrix(config.quantfile());

        printf("...done\n");

        filereader.openFile(config.rawfile());

        filewriter.openFile(config.encfile());

		string fileHeader = filereader.readLine();
		
		stringstream ss(fileHeader);

		string sIntro, sWidth, sHeight; 

		getline(ss, sIntro, ' ');
		getline(ss, sWidth, ' ');
		getline(ss, sHeight, ' ');

		sWidth.erase(0,1);
		sHeight.erase(0,1);

		config.setWidth(atoi(sWidth.data()));
		config.setHeight(atoi(sHeight.data()));

		printf("Initializing encoder...\n");

        Encoder encoder(config, qm, 4);
        Encoder encoderPFrame(config, qm, 16);

        printf("...done\n");
		printf("Beginning encoding...\n");

        filewriter.writeFileHeader(config.getWidth(), config.getHeight(), *qm);

        uint8_t* previousFrame = NULL;

        int gopSize = config.getGOP();

        int gopPosition = 0;

        Decoder decoderIFrame(config.getWidth(), config.getHeight(), qm, 4);
        Decoder decoderPFrame(config.getWidth(), config.getHeight(), qm, 16);

		int totalFrameslength = (int)(filereader.getFileSize() / (config.getWidth()*config.getHeight()*1.5f));

		int frameCounter = 0;

		size_t size = config.getWidth()*config.getHeight();

		string frameIndicator;

        while((frameIndicator = filereader.readLine()).compare("EOF") != 0)
        {
			buff = filereader.readUncompressedData(size, success);
            util::BitStreamWriter* bsw;
            int type = 0;
            if(gopPosition == 0)
            {
				NsightDebug::pushTracerStack("IFrame");
#if I_USE_CUDA == 1
                bsw = encoder.encode(buff->data(), true);
#else
				bsw = encoder.encode(buff->data(), false);
#endif
                util::BitStreamReader bsr(bsw->get_buffer(), bsw->get_size());
				if(previousFrame != 0)
					delete[] previousFrame;
                previousFrame = decoderIFrame.decode(bsr);
                type = 0;
            }
            else
            {
				NsightDebug::pushTracerStack("PFrame");
				NsightDebug::pushTracerStack("Phase 1");
				uint8_t* temp = new uint8_t[size];
				std::memcpy(temp, buff->data(), size * sizeof(uint8_t));

                bsw = encoder.encode(buff->data(), false);
                util::BitStreamReader bsr(bsw->get_buffer(), bsw->get_size());
                uint8_t* pFrame = decoderIFrame.decode(bsr);
				delete bsw->get_buffer();
                delete bsw;
				NsightDebug::pop();
				NsightDebug::pushTracerStack("Phase 2");
#if P_USE_CUDA == 1
                bsw = encoderPFrame.encodePFrame(temp, pFrame, previousFrame, true);
#else
				bsw = encoderPFrame.encodePFrame(temp, pFrame, previousFrame, false);
#endif
				NsightDebug::pop();
                util::BitStreamReader bsr2(bsw->get_buffer(), bsw->get_size());
				if(previousFrame != 0)
					delete[] previousFrame;
                previousFrame = pFrame;
                type = 1;
				delete[] temp;
            }

			NsightDebug::pop();

            uint32_t bytesWritten = (uint32_t)(bsw->get_position()/8);
            if(bsw->get_position()%8 != 0)
                ++bytesWritten;
            filewriter.writeFrameHeader(type, bytesWritten);
            filewriter.writeData(bsw);
			delete bsw->get_buffer();
            delete bsw;
            gopPosition++;
            if(gopPosition >= gopSize)
                gopPosition = 0;

			++frameCounter;
			delete buff;

			std::cout << "Progress: " << (int)((float)frameCounter / (float)totalFrameslength * 100) << "%\r";
        }

		if(previousFrame != NULL)
			delete[] previousFrame;

		delete qm;

        filereader.closeFile();

        filewriter.closeFile();

        printf("...done\n");

		nvtxRangePop();

        return 0;

    }
}
